package com.gitee.wsl.logger.context

import com.gitee.wsl.platform.thread.ThreadLocalRef
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.currentCoroutineContext
import kotlinx.coroutines.withContext
import kotlin.coroutines.AbstractCoroutineContextElement
import kotlin.coroutines.CoroutineContext


interface LogContext : Map<String, Any> {

    operator fun plus(other: LogContext?): LogContext

    fun clone(): LogContext = LogContext(toMap())

    fun toImmutable(): LogContext = LogContext(toMap())

    fun toMutable(): MutableLogContext = MutableLogContext(toMutableMap())
}

operator fun LogContext.plus(other: Map<String, Any>?): LogContext = if (other == null) this else {
    LogContext((this as Map<String, Any>).plus(other))
}

internal class LogContextImpl(data: Map<String, Any>) : LogContext, Map<String, Any> by data {

    override fun plus(other: LogContext?): LogContext = if (other == null) this else {
        LogContext((this as Map<String, Any>).plus(other))
    }
}

fun LogContext(data: Map<String, Any>): LogContext = LogContextImpl(data)

object EmptyLogContext : LogContext by LogContext(emptyMap())

interface MutableLogContext : LogContext, MutableMap<String, Any>,CoroutineContext.Element {

    /**
     * Key used to retrieve the [LogContext] from the coroutine context.
     */
     companion object Key : CoroutineContext.Key<MutableLogContext>

    override operator fun plus(other: LogContext?): MutableLogContext = if (other == null) clone() else {
        val data = mutableMapOf<String, Any>()

        data.putAll(this)
        data.putAll(other)

        MutableLogContext(data)
    }

    operator fun plusAssign(other: LogContext?) {
        if (other != null) {
            putAll(other)
        }
    }

    fun snap() = LogContext(toMap())

    override fun clone(): MutableLogContext = MutableLogContext(toMutableMap())

    /**
     * Put zero or more items into the [LogContext].
     * @param newItems context items to add
     */
     fun putItems(vararg newItems: ContextItem)

    /**
     * Remove zero or more items from the [LogContext].
     * @param keys keys of items to remove
     */
     fun removeItem(vararg keys: String)

}

internal class MutableLogContextImpl(
    data: MutableMap<String, Any>,
) : MutableLogContext,
    MutableMap<String, Any> by HashMap(),
    AbstractCoroutineContextElement(MutableLogContext)
{

    init {
        putAll(data)
    }

    /**
     * Put zero or more items into the [LogContext].
     * @param newItems context items to add
     */
    override fun putItems(vararg newItems: ContextItem) {
        newItems.forEach {
            put(it.first,it.second)
        }
    }

    /**
     * Remove zero or more items from the [LogContext].
     * @param keys keys of items to remove
     */
    override fun removeItem(vararg keys: String) {
        keys.forEach { remove(it) }
    }

}

fun MutableLogContext(data: MutableMap<String, Any> = mutableMapOf()): MutableLogContext = MutableLogContextImpl(data)

object GlobalLogContext : MutableLogContext by MutableLogContext()

class ThreadLogContext : MutableLogContext ,
    AbstractCoroutineContextElement(MutableLogContext){

    private val threadLocal = ThreadLocalRef{ MutableLogContext() }

    private val data: MutableLogContext
        get() {
            var data: MutableLogContext? = threadLocal.get()

            if (data == null) {
                data = MutableLogContext()
                threadLocal.set(data)
            }

            return data
        }

    override val entries: MutableSet<MutableMap.MutableEntry<String, Any>>
        get() = data.entries

    override val keys: MutableSet<String>
        get() = data.keys

    override val size: Int
        get() = data.size

    override val values: MutableCollection<Any>
        get() = data.values

    override fun containsKey(key: String): Boolean = data.containsKey(key)

    override fun containsValue(value: Any): Boolean = data.containsValue(value)

    override fun get(key: String): Any? = data[key]

    override fun isEmpty(): Boolean = data.isEmpty()

    override fun clear() = data.clear()

    override fun put(key: String, value: Any): Any? = data.put(key, value)

    override fun putAll(from: Map<out String, Any>) = data.putAll(from)

    override fun remove(key: String): Any? = data.remove(key)

    /**
     * Put zero or more items into the [LogContext].
     * @param newItems context items to add
     */
    override fun putItems(vararg newItems: ContextItem) {
        newItems.forEach {
            put(it.first,it.second)
        }
    }

    /**
     * Remove zero or more items from the [LogContext].
     * @param keys keys of items to remove
     */
    override fun removeItem(vararg keys: String) {
        keys.forEach { remove(it) }
    }
}


/**
 * Utility function that stores items into a [LogContext] in the current
 * coroutine scope.
 *
 * @param items context items to include in the context
 * @param block lambda that runs in the scope
 * @return the value returned by [block]
 */
suspend fun <R> withLogContext(
    vararg items: ContextItem,
    block: suspend CoroutineScope.() -> R,
): R =
    withContext(logContext(*items)) {
        block()
    }

/**
 * Utility function for constructing a [LogContext] with specified items.
 * @param items context items to include in the context
 * @return a [LogContext] containing the specified items
 */
suspend fun logContext(vararg items: ContextItem): CoroutineContext {
    val allItems = currentCoroutineContext()[MutableLogContext]
                       ?.toMutableMap()
                       ?: mutableMapOf()
    allItems.putAll(items)
    return MutableLogContextImpl(allItems)
}

/** Type alias for an item in a context. */
typealias ContextItem = Pair<String, Any>

/**
 * Adds zero or more items to the [LogContext] in the current coroutine scope.
 * @param items context items to add
 */
suspend fun addToContext(vararg items: ContextItem) {
    currentCoroutineContext()[MutableLogContext]?.putItems(*items)
}

/**
 * Removes zero or more items from the [LogContext] in the current coroutine scope.
 * @param keys keys of the items to remove
 */
suspend fun removeFromContext(vararg keys: String) {
    currentCoroutineContext()[MutableLogContext]?.removeItem(*keys)
}
